﻿#-------------------------------------------------------------------------------
# Name:        TestOls
# Purpose:
#
# Author:      SUPER_紫电  QQ:911344755
#
# Created:     16/08/2022
# Copyright:   (c) SUPER_紫电 2022 All rights reserved.
# Licence:     <Modified BSD License>
#-------------------------------------------------------------------------------

import os
import sys
try:
    sys.path.append(os.path.dirname(__file__) + r'\uses')
except:
    sys.path.append(os.path.dirname(__file__) + r'\uses')

import threading
from time import time, sleep
import pci as pci
import taskfile as taskfile
from taskfile import *

def beep():
    freq = 1193180000 / 440000 # 440Hz
    outp(0x43, 0xB6)
    outp(0x42, int(freq) & 0xFF)
    outp(0x42, int(freq) >> 9)
    sleep(0.1)
    outp(0x61, inp(0x61) | 0x03)
    sleep(1)
    outp(0x61, inp(0x61) & 0xFC)

def enum_pci():
    print('\n[PCI Device]\nBUS DEV FUN VID   DID   SVID  SDID  ClassCode')
    for bus in range(256):
        for dev in range(32):
            for func in range(8):
                address = pci.PciBusDevFunc(bus, dev, func)
                vidValue = ols.ReadPciConfigDword(address, 0x00)
                if (vidValue == 0xFFFFFFFF) or (vidValue == 0):
                    continue
                classValue = ols.ReadPciConfigDword(address, 0x08)
                subClass = ols.ReadPciConfigDword(address, 0x2C)
                headerType = ols.ReadPciConfigByte(address, 0x0E)
                vid = vidValue & 0xFFFF
                did = vidValue >> 16 & 0xFFFF
                svid = subClass & 0xFFFF
                sdid = subClass >> 16 & 0xFFFF
                clc =  classValue >> 8 & 0xFFFFFFFF
                fmt = '%02Xh %02Xh %02Xh %04Xh %04Xh %04Xh %04Xh %06Xh'
                print(fmt % (bus, dev, func, vid, did, svid, sdid, clc))

def enum_stor():
    global ata_controller
    ata_controller.clear()
    print('\n[Mass Storage Controllers]')
    pciDev = ''
    bar = pci.PCICFG_REG_BAR0
    for bus in range(256):
        for dev in range(32):
            for func in range(8):
                address = pci.PciBusDevFunc(bus, dev, func)
                vidValue = ols.ReadPciConfigDword(address, 0x00)
                if (vidValue == 0xFFFFFFFF) or (vidValue == 0):
                    continue
                classValue = ols.ReadPciConfigDword(address, 0x08)
                subClass = ols.ReadPciConfigDword(address, 0x2C)
                headerType = ols.ReadPciConfigByte(address, 0x0E)
                if classValue >> 24 & 0xFF != 1:
                    continue
                vid = vidValue & 0xFFFF
                did = vidValue >> 16 & 0xFFFF
                svid = subClass & 0xFFFF
                sdid = subClass >> 16 & 0xFFFF
                clc =  classValue >> 8 & 0xFFFFFFFF
                pciDev = 'BUS DEV FUN VID   DID   SVID  SDID  ClassCode\n'
                fmt = '%02Xh %02Xh %02Xh %04Xh %04Xh %04Xh %04Xh %06Xh'
                pciDev += fmt % (bus, dev, func, vid, did, svid, sdid, clc)
                print(pciDev)
                print('%04Xh - PciBusDevFunc' % address)
                addressById = ols.FindPciDeviceById(vid, did, 0)
                print('%04Xh - FindPciDeviceById' % addressById)
                baseClass = classValue >> 24 & 0xFF
                subClass = classValue >> 16 & 0xFF
                programIf = classValue >> 8 & 0xFF
                addressByClass = ols.FindPciDeviceByClass(
                    baseClass, subClass, programIf, 0)
                print('%04Xh - FindPciDeviceByClass' % addressByClass)
                #pciconfig = bytearray(256)
                pciCfg = '[PCI Confguration Space]\n'
                pciCfg+='    | 00 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E 0F\n'
                pciCfg+='-----------------------------------------------------\n'
                for i in range(0, 256, 16):
                    pciCfg += '%04X|' % i
                    for j in range(16):
                        data = ols.ReadPciConfigByte(address, i + j)
                        #pciCfgconfig[i + j] = data
                        pciCfg += ' %02X' % data
                    pciCfg += '\n'
                print(pciCfg)
                if headerType != 0:
                    continue
                for i in range(3):
                    base = ols.ReadPciConfigDword(address, bar + i * 8)
                    alt = ols.ReadPciConfigDword(address, bar + i * 8 + 4)
                    if(base&1) and (alt&1) and (base>0x170) and (alt<0xFFFF):
                        ata_port = {}
                        ata_port['base'] = base - 1
                        ata_port['alt'] = alt + 1
                        ata_controller.append(ata_port)
    print('[Mass Storage IO Port]\n Base   Alt')
    for d in ata_controller:
        #print(' '.join(['%s:0x%04X' % (x,y) for x,y in d.items()]))
        print(' '.join(['0x%04X' % x for x in d.values()]))

def get_status():
    print('\n[Driver Status] 0x%04X 0x%04X' % (ata_port['base'], ata_port['alt']))
    status = ols.ReadIoPortByte(ata_port['base'] + 7)
    print('0x%02X' % status)

def get_char(buffer, offset, size):
    temp = ''
    j = 0
    r = int(size/2)
    for i in range(r):
        temp += chr(buffer[offset + j + 1])
        temp += chr(buffer[offset + j])
        j += 2
    return temp.strip()

def get_ident():
    global ata_port
    for ata_port in ata_controller:
        set_reg_addr(ata_port['base'], ata_port['alt'])
        #print(ata_reg)
        get_status()
        if eob:
            soft_reset()
        if eto(10):
            print('Error or time out!')
            continue
        if not cmd_identify():
            continue
        identify = bytearray()
        ident = '\n[Identify Device Data]\n'
        ident += '    | 00 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E 0F |\n'
        #ident += '-------------------------------------------------------'
        #ident += '-----------------\n'
        for i in range(73):
            ident += '-'
        for i in range(0, 512, 16):
            ident += '\n%04X|' % i
            dec = ''
            for j in range(0, 16, 2):
                data = inpw(ata_port['base'])
                high = data & 0xFF
                low = data >> 8
                identify.append(high)
                identify.append(low)
                ident += ' %02X %02X' % (high, low)
                dec += '%c%c' % (high, low)
            ident += ' | %s' % dec.replace('\n', '').replace('\r', '')
        if len(identify) == 512:
            model = get_char(identify, 54, 40)
            fw = get_char(identify, 46, 8)
            serial = get_char(identify, 20, 20)
            filename = model + '_' +  fw + '_' + serial + '.bin'
            with open(filename, 'wb') as file:
                file.write(identify)
        print(ident)

def test_ols():
    funs = [beep, enum_pci, enum_stor, get_ident]
    for fun in funs:
        th = threading.Thread(target=fun)
        th.start()
        th.join()

def test_fun(fun):
    th = threading.Thread(target=fun)
    th.start()
    th.join()

def cls():
    os.system('cls')

dict_function = {
    '1': beep,
    '2': enum_pci,
    '3': enum_stor,
    '4': get_ident,
    '5': test_ols,
    'c': cls,
}

def main():
    if not init_drv():
        print('Driver initialize %s!' % str(bool(taskfile.drv_initialized)).lower())
        return
    print('Driver initialize %s.' % str(bool(taskfile.drv_initialized)).lower())
    test_ols()
    while True:
        print('\n1 - Beep')
        print('2 - PCI Device')
        print('3 - Mass Storage Controllers')
        print('4 - Identify Device Data')
        print('5 - Test All')
        print('c - Clear')
        print('x - Exit')
        selection = input('Selection: ')
        if selection.lower() in dict_function:
            test_fun(dict_function[selection.lower()])
            continue
        if selection.lower() == 'x':
            break
        else:
            os.system('cls')
            continue
    deinit_drv()

if __name__ == '__main__':
    main()
